﻿#include "PBiCGStab.hpp"

template <typename Cmpt>
UNAP::PBiCGStab<Cmpt>::PBiCGStab(Communicator *other_comm)
    : Solver<Cmpt>(other_comm), deletePrecondPtr_(false), precondPtr_(NULL) {
  precondPtr_ = new Preconditioner<Cmpt>(other_comm);
  deletePrecondPtr_ = true;
}

template <typename Cmpt>
UNAP::PBiCGStab<Cmpt>::PBiCGStab(Preconditioner<Cmpt> &precond)
    : Solver<Cmpt>(NULL), deletePrecondPtr_(false), precondPtr_(NULL) {
  precondPtr_ = &precond;
  this->setCommunicator(precond.getCommunicator());
}

template <typename Cmpt>
UNAP::SolverPerformance UNAP::PBiCGStab<Cmpt>::solve(Vector<Cmpt> &x,
                                                     const Matrix<Cmpt> &A,
                                                     const Vector<Cmpt> &b) {
  SolverPerformance solverPerf;
  const label rowSize = x.size();
  Vector<Cmpt> pA(rowSize, this->commcator_);
  Vector<Cmpt> yA(rowSize, this->commcator_);
  Vector<Cmpt> rA(rowSize, this->commcator_);
  Vector<Cmpt> rA0(rowSize, this->commcator_);
  Vector<Cmpt> AyA(rowSize, this->commcator_);
  Vector<Cmpt> sA(rowSize, this->commcator_);
  Vector<Cmpt> zA(rowSize, this->commcator_);
  Vector<Cmpt> tA(rowSize, this->commcator_);
  Cmpt rA0rA = 0., alpha = 0., omega = 0.;
  Cmpt rA0rAold, beta, rA0AyA, tAtA;

  A.spMV(yA, x);
  waxpy(rA, -1.0, yA, b);
  rA0 = rA;
  // adjust init residual
  solverPerf.initialResidual(this->normFactor(b));

#ifdef VERBOSE
  if (this->ifPrint_) solverPerf.printResidual();
#endif

  do {
    solverPerf.nIterations()++;

    rA0rAold = rA0rA;

    rA0rA = dot(rA0, rA);

    if (solverPerf.checkSingularity(ABS(rA0rA))) {
#ifdef VERBOSE
      if (this->ifPrint_) {
        this->commcator_->log() << "singularity! rA0rA = " << rA0rA << ENDL;
      }
#endif
      break;
    }

    if (solverPerf.nIterations() == 1) {
      pA = rA;
    } else {
      beta = (rA0rA / rA0rAold) * (alpha / omega);
      axpbypz(pA, -omega * beta, beta, AyA, rA);
    }

    precondPtr_->precondition(yA, pA);

    A.spMV(AyA, yA);
    rA0AyA = dot(rA0, AyA);
    alpha = rA0rA / rA0AyA;

    waxpy(sA, -alpha, AyA, rA);

    precondPtr_->precondition(zA, sA);

    A.spMV(tA, zA);

    tAtA = tA.SumSqr();

    if (solverPerf.checkSingularity(sqrt(tAtA))) {
      axpy(x, alpha, yA);
      solverPerf.finalResidual(this->normFactor(sA));
#ifdef VERBOSE
      if (this->ifPrint_) solverPerf.printResidual();
#endif
      break;
    }

    omega = dot(tA, sA);

    omega /= tAtA;

    axpy(x, alpha, yA);
    axpy(x, omega, zA);
    waxpy(rA, -omega, tA, sA);

    solverPerf.finalResidual(this->normFactor(rA));
#ifdef VERBOSE
    if (this->ifPrint_) solverPerf.printResidual();
#endif
  } while ((solverPerf.nIterations() < this->maxIter_ &&
            !(solverPerf.checkConvergence(this->tolerance_, this->relTol_,
                                          this->minIter_))));

  return solverPerf;
}

template class UNAP::PBiCGStab<scalar>;
